{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OchreLang.TopSyntax
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Top level syntax.
--
--------------------------------------------------------------------------------

module OchreLang.TopSyntax
  (

    Program(..)
  , ModuleDict
  , Module(..)
  , Instr(..)
  , Trait(..)
  , TraitDefn(..)
  , TraitExpr(..)
  , TraitAlias(..)
  , MethodDecl(..)
  , MethodDefn(..)
  , Decl(..)
  , DoBlock
  , Stmt(..)
  , VarBind(..)
  , Expr(..)
  , Literal(..)

  -- * Traversal
  , mapM_Module

  -- * Pretty print
  , ppProgram
  , ppModule
  , ppInstr
  , ppTrait
  , ppIdent
  , ppDecl
  , ppExpr
  , ppStmt
  , ppLiteral


  ) where


import qualified OchreLang.OrchSpec as OS
import OchreLang.SyntaxBase
import OchreLang.Utils

import Text.PrettyExprHPJ                       -- package: pretty-expr-hpj
import Text.PrettyExprHPJ.Haskell98

import Text.PrettyPrint.HughesPJ                -- package: pretty


import qualified Data.Map as Map
import qualified Data.Traversable as T

--------------------------------------------------------------------------------
-- 


-- Note - syntax is not a one-to-one representation of files.
-- The Main module should have an OrchSpec which is floated out
-- in the concrete syntax to Program level.


data Program = Program 
    { program_orch_spec         :: OS.OrchSpec
    , program_modules           :: ModuleDict
    , program_main_name         :: String
    }
  deriving (Eq,Show)

type ModuleDict = Map.Map ModuleName Module

data Module = Module
    { module_name               :: ModuleName
    , module_imports            :: [ModuleName]
    , module_instrs             :: [Instr]
    , module_traits             :: [Trait]
    }
  deriving (Eq,Show)



data Instr = Instr 
    { instr_name                :: Ident
    , instr_superclass          :: Ident
    , instr_opt_trait           :: Maybe TraitExpr
    , instr_methods             :: [MethodDecl]
    }
  deriving (Eq,Show)

data MethodDecl = MethodDecl Ident MethodDefn
  deriving (Eq,Show)

data MethodDefn = MethodDefn [Ident] DoBlock 
  deriving (Eq,Show)

data Trait = Trait Ident TraitDefn 
  deriving (Eq,Show)

data TraitDefn = TraitDefn (Maybe TraitExpr) [MethodDecl]
  deriving (Eq,Show)





data TraitExpr = NamedTrait   Ident
               | TraitComp    TraitExpr   TraitExpr
               | TraitExclude TraitExpr   [Ident]
               | TraitRename  TraitExpr   [TraitAlias]
  deriving (Eq,Show)

data TraitAlias = TraitAlias Ident Ident
  deriving (Eq,Show)
  
  
-- See H98 Report 3.14
-- Must be able to nest dos

type DoBlock = [Stmt]


-- | In TopSyntax lambdas are /statements/ rather than 
-- expressions.
--
-- This is to keep the separation between do-stmts and pure 
-- expressions.
--
data Stmt = Return    Expr
          | PrimCall  Ident     [Expr]
          | DoBind    VarBind   DoBlock
          | DoLet     Decl
          | DoLam     Ident     [Ident]   DoBlock
          | DoIf      Expr      DoBlock   DoBlock
  deriving (Eq,Show)



data VarBind = Bind1     Ident
             | BindTuple [Ident]
  deriving (Eq,Show)

-- | Decls only support binding simple values (cannot bind 
-- functions).
--
data Decl = Decl Ident Expr
  deriving (Eq,Show)


-- | Expressions are pure (?)
--
-- Applications are builtin Csound (pure) functions like abs.
--
-- There are no pfields in the top level syntax - pfields are
-- positional and hence not modular. We synthesize them from the 
-- orchestra spec later in compilation.
-- 
data Expr = Lit     Literal
          | Var     Ident
          | App     Ident     [Expr]
          | Cond    Expr      Expr      Expr
          | UnaryE  UnaryOp   Expr
          | BinE    BinOp     Expr      Expr
          | RelE    RelOp     Expr      Expr
          | Tuple   [Expr]
  deriving (Eq,Show)




data Literal = Bool     Bool
             | Int      Int
             | Float    Decimal
             | String   String
  deriving (Eq,Show)


--------------------------------------------------------------------------------
-- Traversal


mapM_Module :: Monad m => (Module -> m Module) -> Program -> m Program
mapM_Module mf prgm@(Program { program_modules = moduls }) =
   do { dict <- T.mapM mf moduls
      ; return $ prgm { program_modules = dict }
      }




--------------------------------------------------------------------------------
-- Pretty printing



-- TODO - pretty printing should be improved to @inject@
-- ospec back into main. Round-tripping would be nice for 
-- testing the parser...
--
ppProgram :: Program -> Doc
ppProgram (Program { program_orch_spec = ospec
                   , program_modules   = modules }) = 
    vspaced $ OS.ppOrchSpec ospec : map ppModule module_list
  where
    module_list = Map.elems modules
    


ppModule :: Module -> Doc
ppModule (Module { module_name      = name
                 , module_imports   = imports
                 , module_instrs    = instrs
                 , module_traits    = traits }) = 
    let modu  = text "module" <+> ppModuleName name <+> text "where"
        imps  = vspaced $ map prettyImport imports 
        idocs = map ppInstr instrs
        tdocs = map ppTrait traits
    in vspaced $ modu : imps : (idocs ++ tdocs)
    


prettyImport :: ModuleName -> Doc
prettyImport name = text "import" <+> ppModuleName name

ppInstr :: Instr -> Doc
ppInstr (Instr vid super texp xs) = 
    text "instr" <+> ppIdent vid <+> text "extends" <+> ppIdent super
      $+$ nest 2 (vsemiBraces body)
  where
    body = ppUsesDecl texp : map ppMethodDecl xs

ppUsesDecl :: Maybe TraitExpr -> Doc
ppUsesDecl Nothing  = text "uses" <+> braces empty
ppUsesDecl (Just e) = text "uses" <+> braces (ppTraitExpr e)

ppTrait :: Trait -> Doc
ppTrait (Trait vid body) = 
    text "trait" <+> ppIdent vid $+$ nest 2 (ppTraitDefn body)


ppTraitDefn :: TraitDefn -> Doc
ppTraitDefn (TraitDefn texp xs) = vsemiBraces body
  where
    body = ppUsesDecl texp : map ppMethodDecl xs


ppTraitExpr :: TraitExpr -> Doc
ppTraitExpr (NamedTrait vid)    = ppIdent vid
ppTraitExpr (TraitComp e1 e2)   = 
    ppTraitExpr e1 <+> char '&' <+> ppTraitExpr e2

ppTraitExpr (TraitExclude e xs) = 
    ppTraitExpr e <+> text "minus" <+> commaBraces (map ppIdent xs)

ppTraitExpr (TraitRename e xs)  = 
    ppTraitExpr e <> text "with" <+> semiBraces (map ppTraitAlias xs)

ppTraitAlias :: TraitAlias -> Doc
ppTraitAlias (TraitAlias x y) = ppIdent x <+> text "->" <+> ppIdent y


ppMethodDecl :: MethodDecl -> Doc
ppMethodDecl (MethodDecl vid body) = 
    ppIdent vid <+> char '=' <+> ppMethodDefn body


ppMethodDefn :: MethodDefn -> Doc
ppMethodDefn (MethodDefn args body) 
    | null args = ppDoBlock body
    | otherwise = let formals = hsep $ map ppIdent args
                  in char '\\' <> formals <+> text "->" <+> ppDoBlock body

ppDecl :: Decl -> Doc
ppDecl (Decl vid e)             = ppIdent vid <+> char '=' <+> ppExpr e


ppDoBlock :: DoBlock -> Doc
ppDoBlock [x]                   = ppStmt x
ppDoBlock xs                    = text "do" <+> vsemiBraces (map ppStmt xs)

ppStmt :: Stmt -> Doc
ppStmt (Return e)               = text "return" <+> ppExpr e
ppStmt (PrimCall v es)          = ppIdent v <+> commaSep (map ppExpr es)
ppStmt (DoBind a de)            = ppVarBind a <+> text "<-" <+> ppDoBlock de
ppStmt (DoLet d)                = text "let" <+> ppDecl d
ppStmt (DoLam s args de)        = 
    text "letfun" <+> ppIdent s <+> hsep (map ppIdent args) 
                  <+> char '=' <+> ppDoBlock de 
ppStmt (DoIf c d1 d2)           = 
    text "if" <+> ppExpr c <+> nest 2 (    text "then" <+> ppDoBlock d1
                                       $+$ text "else" <+> ppDoBlock d2)



ppVarBind :: VarBind -> Doc
ppVarBind (Bind1 v)             = ppIdent v
ppVarBind (BindTuple vs)        = parens (commaSep $ map ppIdent vs)

ppExpr :: Expr -> Doc
ppExpr = unparse . exprDoc

exprDoc :: Expr -> DocE
exprDoc (Lit v)                 = Atom $ ppLiteral v
exprDoc (Var s)                 = Atom $ ppIdent s
exprDoc (App vid es)            = Atom $ ppIdent vid <> parens (commaSep $ map ppExpr es)
exprDoc (Cond c e1 e2)          = 
    Atom $ text "cond" <+> ppExpr c <+> char ':' <+> ppExpr e1
                                    <+> char '?' <+> ppExpr e2

exprDoc (UnaryE op e)           = unaryDocE op (exprDoc e)
exprDoc (BinE op e1 e2)         = binDocE op (exprDoc e1) (exprDoc e2)
exprDoc (RelE relop e1 e2)      = relDocE relop (exprDoc e1) (exprDoc e2)
exprDoc (Tuple es)              = Atom $ tupled (map ppExpr es) 


ppLiteral :: Literal -> Doc
ppLiteral (Bool b)              = text $ show b
ppLiteral (Int i)               = int i
ppLiteral (Float d)             = text $ truncatedDouble $ realToFrac d
ppLiteral (String s)            = doubleQuotes $ text s

